#include <mcp/transport/websocket_transport.h>
#include <mcp/error.h>
#include <mcp/json/value.h>

#include <iostream>
#include <sstream>
#include <chrono>
#include <memory>
#include <thread>
#include <vector>
#include <map>

namespace mcp {
namespace transport {

// WebSocket client simulation for demonstration
class WebSocketClientSim {
public:
    struct Message {
        std::string data;
        bool is_binary = false;
    };

    bool connect(const std::string& url) {
        std::cout << "[WebSocket] Connecting to: " << url << std::endl;
        return true;
    }

    void disconnect() {
        std::cout << "[WebSocket] Disconnecting" << std::endl;
    }

    bool send(const std::string& data, bool is_binary = false) {
        std::cout << "[WebSocket] Sending: " << data.substr(0, 100) << "..." << std::endl;
        return true;
    }

    bool send_ping() {
        std::cout << "[WebSocket] Sending ping" << std::endl;
        return true;
    }

    bool poll(Message& message) {
        // Simulate receiving messages
        static int counter = 0;
        if (counter++ % 10 == 0) {
            message.data = "{\"jsonrpc\":\"2.0\",\"method\":\"test\",\"id\":" + std::to_string(counter) + "}";
            message.is_binary = false;
            return true;
        }
        return false;
    }
};

// WebSocketTransport implementation
struct WebSocketTransport::Impl {
    WebSocketConfig config;
    std::atomic<bool> running{false};
    std::atomic<bool> connected{false};
    std::unique_ptr<std::thread> receiver_thread;
    std::unique_ptr<std::thread> ping_thread;
    WebSocketClientSim client;
    MessageHandler message_handler;
    ErrorHandler error_handler;
    CloseHandler close_handler;
    OpenHandler open_handler;
    mutable std::mutex transport_mutex;

    Impl(const WebSocketConfig& cfg) : config(cfg) {}

    void start_receiving();
    void stop_receiving();
    void start_pinging();
    void stop_pinging();
    void handle_messages();
};

WebSocketTransport::WebSocketTransport(const WebSocketConfig& config) : impl_(std::make_unique<Impl>(config)) {}

WebSocketTransport::WebSocketTransport(uint16_t port)
    : impl_(std::make_unique<Impl>(WebSocketConfig{.port = port})) {}

WebSocketTransport::~WebSocketTransport() {
    stop();
}

bool WebSocketTransport::start() {
    std::lock_guard<std::mutex> lock(impl_->transport_mutex);

    if (impl_->running) {
        return true;
    }

    // Connect to WebSocket server
    std::string url = get_url();
    if (!impl_->client.connect(url)) {
        if (impl_->error_handler) {
            impl_->error_handler(Error::connection_error("Failed to connect to WebSocket server"));
        }
        return false;
    }

    impl_->running = true;
    impl_->connected = true;

    // Start the receiver thread
    impl_->receiver_thread = std::make_unique<std::thread>([this]() {
        impl_->start_receiving();
    });

    // Start the ping thread if enabled
    if (impl_->config.ping_interval_seconds > 0) {
        impl_->ping_thread = std::make_unique<std::thread>([this]() {
            impl_->start_pinging();
        });
    }

    // Notify open handler
    if (impl_->open_handler) {
        impl_->open_handler();
    }

    return true;
}

void WebSocketTransport::stop() {
    std::lock_guard<std::mutex> lock(impl_->transport_mutex);

    if (!impl_->running) {
        return;
    }

    impl_->running = false;
    impl_->connected = false;

    // Stop the ping thread
    impl_->stop_pinging();

    // Stop the receiver thread
    impl_->stop_receiving();

    // Disconnect from WebSocket server
    impl_->client.disconnect();

    if (impl_->receiver_thread && impl_->receiver_thread->joinable()) {
        impl_->receiver_thread->join();
        impl_->receiver_thread.reset();
    }

    if (impl_->ping_thread && impl_->ping_thread->joinable()) {
        impl_->ping_thread->join();
        impl_->ping_thread.reset();
    }

    // Notify close handler
    if (impl_->close_handler) {
        impl_->close_handler();
    }
}

bool WebSocketTransport::is_connected() const {
    return impl_->connected;
}

bool WebSocketTransport::send(const json::Value& message) {
    if (!impl_->running || !impl_->connected) {
        return false;
    }

    try {
        std::string serialized = message.serialize();

        // Check message size limit
        if (serialized.size() > impl_->config.max_message_size) {
            if (impl_->error_handler) {
                impl_->error_handler(Error::validation_error("Message too large"));
            }
            return false;
        }

        bool success = impl_->client.send(serialized, false);
        if (!success) {
            if (impl_->error_handler) {
                impl_->error_handler(Error::network_error("Failed to send WebSocket message"));
            }
            return false;
        }

        return true;
    } catch (const std::exception& e) {
        if (impl_->error_handler) {
            impl_->error_handler(Error::serialization_error(e.what()));
        }
        return false;
    }
}

std::coroutine_handle<> WebSocketTransport::send_async(const json::Value& message) {
    // For WebSocket transport simulation, send_async behaves the same as send
    // In a real implementation, this would use async WebSocket client
    send(message);
    return std::noop_coroutine();
}

void WebSocketTransport::set_message_handler(MessageHandler handler) {
    impl_->message_handler = std::move(handler);
}

void WebSocketTransport::set_error_handler(ErrorHandler handler) {
    impl_->error_handler = std::move(handler);
}

void WebSocketTransport::set_close_handler(CloseHandler handler) {
    impl_->close_handler = std::move(handler);
}

void WebSocketTransport::set_open_handler(OpenHandler handler) {
    impl_->open_handler = std::move(handler);
}

std::string WebSocketTransport::get_type() const {
    return "websocket";
}

json::Value WebSocketTransport::get_capabilities() const {
    json::Value capabilities = json::Value::object();
    capabilities["type"] = "websocket";
    capabilities["supports_async"] = true;
    capabilities["supports_compression"] = impl_->config.enable_compression;
    capabilities["supports_ping_pong"] = true;
    capabilities["subprotocol"] = impl_->config.subprotocol;
    capabilities["max_message_size"] = static_cast<int64_t>(impl_->config.max_message_size);
    capabilities["ping_interval"] = impl_->config.ping_interval_seconds;
    capabilities["pong_timeout"] = impl_->config.pong_timeout_seconds;
    return capabilities;
}

std::string WebSocketTransport::get_url() const {
    std::ostringstream ss;
    ss << "ws://" << impl_->config.host << ":" << impl_->config.port << impl_->config.path;
    return ss.str();
}

uint16_t WebSocketTransport::get_port() const {
    return impl_->config.port;
}

bool WebSocketTransport::send_ping() {
    if (!impl_->running || !impl_->connected) {
        return false;
    }

    return impl_->client.send_ping();
}

void WebSocketTransport::update_config(const WebSocketConfig& config) {
    std::lock_guard<std::mutex> lock(impl_->transport_mutex);
    impl_->config = config;
}

const WebSocketConfig& WebSocketTransport::get_config() const {
    return impl_->config;
}

void WebSocketTransport::Impl::start_receiving() {
    while (running) {
        try {
            handle_messages();
            std::this_thread::sleep_for(std::chrono::milliseconds(50));
        } catch (const std::exception& e) {
            if (error_handler) {
                error_handler(Error::network_error(e.what()));
            }
        }
    }
}

void WebSocketTransport::Impl::stop_receiving() {
    // No specific action needed for simulation
    // The loop in start_receiving will exit when running is false
}

void WebSocketTransport::Impl::start_pinging() {
    while (running && config.ping_interval_seconds > 0) {
        try {
            std::this_thread::sleep_for(std::chrono::seconds(config.ping_interval_seconds));

            if (running && connected) {
                if (!client.send_ping()) {
                    // Ping failed, connection might be lost
                    if (error_handler) {
                        error_handler(Error::connection_error("WebSocket ping failed"));
                    }
                    break;
                }
            }
        } catch (const std::exception& e) {
            if (error_handler) {
                error_handler(Error::network_error(e.what()));
            }
        }
    }
}

void WebSocketTransport::Impl::stop_pinging() {
    // No specific action needed for simulation
    // The loop in start_pinging will exit when running is false
}

void WebSocketTransport::Impl::handle_messages() {
    WebSocketClientSim::Message message;

    while (client.poll(message)) {
        try {
            // Parse JSON message
            auto result = json::Value::parse(message.data);
            if (result && message_handler) {
                message_handler(*result);
            } else if (!result && error_handler) {
                error_handler(result.error());
            }
        } catch (const std::exception& e) {
            if (error_handler) {
                error_handler(Error::parse_error(e.what()));
            }
        }
    }
}

// WebSocketTransportFactory implementation
std::unique_ptr<Transport> WebSocketTransportFactory::create() {
    return std::make_unique<WebSocketTransport>();
}

std::unique_ptr<Transport> WebSocketTransportFactory::create_with_config(const WebSocketConfig& config) {
    return std::make_unique<WebSocketTransport>(config);
}

std::unique_ptr<Transport> WebSocketTransportFactory::create_with_port(uint16_t port) {
    return std::make_unique<WebSocketTransport>(port);
}

std::string WebSocketTransportFactory::get_type() const {
    return "websocket";
}

std::vector<std::string> WebSocketTransportFactory::get_supported_schemes() const {
    return {"ws", "wss"};
}

// Utility functions for WebSocket transport
std::unique_ptr<WebSocketTransport> create_websocket_transport(uint16_t port) {
    return std::make_unique<WebSocketTransport>(port);
}

std::unique_ptr<WebSocketTransport> create_websocket_transport_with_config(const WebSocketConfig& config) {
    return std::make_unique<WebSocketTransport>(config);
}

WebSocketConfig create_websocket_config(const std::string& host, uint16_t port, const std::string& path) {
    WebSocketConfig config;
    config.host = host;
    config.port = port;
    config.path = path;
    return config;
}

} // namespace transport
} // namespace mcp